# Requires: PowerShell 5+
# Purpose: One-click install of CUDA-enabled PyTorch into a conda env, then verify GPU availability.
# Default env: video_sr (change via -EnvName)

param(
    [string]$EnvName = 'video_sr',
    [string]$CudaIndexUrl = 'https://download.pytorch.org/whl/cu121',
    [string]$CondaRoot = $null,
    [switch]$RunTest = $true
)

$ErrorActionPreference = 'Stop'
[Console]::OutputEncoding = [System.Text.Encoding]::UTF8

function Write-Info($msg) { Write-Host "[INFO] $msg" -ForegroundColor Cyan }
function Write-Ok($msg) { Write-Host "[OK]   $msg" -ForegroundColor Green }
function Write-Warn($msg) { Write-Host "[WARN] $msg" -ForegroundColor Yellow }
function Write-Err($msg) { Write-Host "[ERR]  $msg" -ForegroundColor Red }

function Get-CondaInfoJson {
    try {
        $out = conda info --json | Out-String
        return ($out | ConvertFrom-Json)
    } catch {
        return $null
    }
}

function Resolve-CondaRoot([string]$Provided) {
    if ($Provided -and (Test-Path $Provided)) { return $Provided }
    $ci = Get-CondaInfoJson
    if ($ci -and $ci.root_prefix -and (Test-Path $ci.root_prefix)) { return $ci.root_prefix }
    # Try to infer from conda.exe path
    try {
        $condaCmd = Get-Command conda -ErrorAction Stop
        # conda.exe is typically at <root>\Scripts\conda.exe
        $scriptsDir = Split-Path $condaCmd.Path -Parent
        $root = Split-Path $scriptsDir -Parent
        if (Test-Path $root) { return $root }
    } catch {}
    # Last resort: common default install path
    $fallbacks = @(
        'C:\ProgramData\Anaconda3',
        'C:\Users\Public\Anaconda3',
        'D:\SoftwareSetUp\anaconda3',
        'C:\Users\'+$env:USERNAME+'\Anaconda3'
    )
    foreach ($p in $fallbacks) { if (Test-Path $p) { return $p } }
    throw "Unable to resolve Conda root. Ensure 'conda' is on PATH or pass -CondaRoot."
}

function Get-EnvPython([string]$root, [string]$envName) {
    $envPath = Join-Path $root ("envs\$envName")
    $py = Join-Path $envPath 'python.exe'
    if (-not (Test-Path $py)) { throw "Python not found at: $py. Verify env name '$envName'." }
    return $py
}

function Cleanup-InvalidPipArtifacts([string]$envPath) {
    $site = Join-Path $envPath 'Lib\site-packages'
    if (-not (Test-Path $site)) { return }
    $targets = @('-ip', '-ip.dist-info')
    foreach ($t in $targets) {
        $p = Join-Path $site $t
        if (Test-Path $p) {
            Write-Warn "Removing invalid pip artifact: $p"
            try { Remove-Item -LiteralPath $p -Recurse -Force -ErrorAction SilentlyContinue } catch {}
        }
    }
}

function Install-PyTorchCuda([string]$pythonExe, [string]$indexUrl) {
    Write-Info "Installing CUDA PyTorch wheels from: $indexUrl"
    $args = @('-m','pip','install','--no-cache-dir','--default-timeout','180','torch','torchvision','torchaudio','--index-url', $indexUrl)
    & $pythonExe $args
}

function Uninstall-PyTorch([string]$pythonExe) {
    Write-Info 'Uninstalling any existing torch/vision/audio packages (pip)'
    & $pythonExe -m pip uninstall -y torch torchvision torchaudio
}

function Conda-Install-PyTorch([string]$envName) {
    Write-Info 'Falling back to conda install (CUDA 12.1)'
    conda install -n $envName pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia -y
}

function Verify-Cuda([string]$pythonExe) {
    Write-Info 'Verifying PyTorch CUDA availability'
    $code = @"
import torch, platform
print('python:', platform.python_version())
print('torch:', torch.__version__, 'cuda:', torch.version.cuda)
print('is_available:', torch.cuda.is_available())
print('device_count:', torch.cuda.device_count())
print('device:', (torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'))
"@
    & $pythonExe -c $code
}

try {
    Write-Info "Resolving Conda root and environment python"
    $root = Resolve-CondaRoot -Provided $CondaRoot
    $envPath = Join-Path $root ("envs\$EnvName")
    $pythonExe = Get-EnvPython -root $root -envName $EnvName
    Write-Ok "Conda root: $root"
    Write-Ok "Env path:  $envPath"
    Write-Ok "Python:     $pythonExe"

    Cleanup-InvalidPipArtifacts -envPath $envPath

    Uninstall-PyTorch -pythonExe $pythonExe

    $installSucceeded = $true
    try {
        Install-PyTorchCuda -pythonExe $pythonExe -indexUrl $CudaIndexUrl
    } catch {
        Write-Warn "pip install encountered an error: $_"
        $installSucceeded = $false
    }

    if (-not $installSucceeded) {
        try {
            Conda-Install-PyTorch -envName $EnvName
            $installSucceeded = $true
        } catch {
            Write-Err "conda install also failed: $_"
            throw
        }
    }

    Verify-Cuda -pythonExe $pythonExe

    # Optionally run test_environment.py if present one level above this script
    $repoRoot = Split-Path $PSScriptRoot -Parent
    $testPy = Join-Path $repoRoot 'test_environment.py'
    if ($RunTest -and (Test-Path $testPy)) {
        Write-Info "Running test script: $testPy"
        & $pythonExe $testPy
    } else {
        Write-Warn "Test script not found or RunTest disabled: $testPy"
    }

    Write-Ok 'Setup complete.'
    Write-Ok 'Tip: Always run with the environment python.exe to avoid base interpreter.'
    Write-Host "Example: `"$pythonExe`" `"$testPy`"" -ForegroundColor Magenta
} catch {
    Write-Err "Setup failed: $_"
    exit 1
}